# main.py

import os
import re
import datetime
import time  # 1. 引入 time 模块
from pysql_compiler.compiler import Compiler
from pysql_compiler.catalog import SemanticError, Catalog
from storage.file_manager import FileManager
from storage.buffer import BufferManager, LRUReplacementPolicy
from engine.executor import Executor
from engine.catalog_manager import CatalogManager
from engine.storage_engine import StorageEngine

def execute_query(sql_query, compiler, executor, catalog_manager, log_file_path):
    try:
        with open(log_file_path, 'a', encoding='utf-8') as log_file:
            timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            log_file.write(f"[{timestamp}] CMD: {sql_query.replace(chr(10), ' ')}\n")
    except IOError as e:
        print(f"错误：无法写入日志文件 {log_file_path}: {e}")

    try:
        # --- 2. 添加计时开始 ---
        start_time = time.time()
        print("--- 正在编译 ---")
        plan, tokens = compiler.compile(sql_query)
        print("编译成功。")

        print("\n--- 正在执行 ---")
        results = []
        plans_to_execute = [plan] if not isinstance(plan, list) else plan
        for p in plans_to_execute:
            results = executor.execute(p)

        # --- 添加计时结束 ---
        end_time = time.time()
        duration = end_time - start_time
        print(f"\n--- 执行完毕，耗时 {duration:.4f} 秒 ---")

        if results and isinstance(results, list):
            header = []
            main_plan_for_header = plan[0] if isinstance(plan, list) else plan

            project_plan = None

            if main_plan_for_header.get('operation') == 'SHOW_TABLES':
                header = ['TABLES']
            else:
                # 否则，使用原有逻辑查找PROJECT节点来确定表头
                project_plan = find_project_node(main_plan_for_header)
                if project_plan:
                    header = extract_header(project_plan, catalog_manager)

            column_types = {}
            if project_plan:
                # 尝试从执行计划中找到源表并获取其schema
                source_node = project_plan.get('source')
                table_schema = None
                while source_node:
                    if source_node.get('table_name'):
                        table_name = source_node.get('table_name')
                        table_schema = catalog_manager.get_table_schema(table_name)
                        break
                    source_node = source_node.get('source')

                if table_schema:
                    column_types = table_schema.columns  # {'col_name': 'TYPE', ...}

            if header:
                print("\n" + " | ".join(map(str, header)))
                print("-" * (sum(len(str(h)) for h in header) + len(header) * 3 - 1))

            for row in results:
                formatted_row = []
                for i, item in enumerate(row):
                    # 通过索引获取列名，并查找其类型
                    col_name = header[i] if i < len(header) else None
                    col_type = column_types.get(col_name, '').upper()

                    if isinstance(item, float):
                        if col_type == 'FLOAT':
                            formatted_row.append(f"{item:.4f}")
                        elif col_type == 'DOUBLE':
                            formatted_row.append(f"{item:.8f}")
                        else:
                            # 如果类型未知（例如计算结果），则默认格式
                            formatted_row.append(str(item))
                    else:
                        formatted_row.append(str(item))
                print(" | ".join(formatted_row))

    except (RuntimeError, ValueError, SemanticError) as e:
        # 将异常向上抛出，以便SOURCE循环可以捕获它
        print(f"发生错误: {e}")
        raise e


def find_project_node(p):
    if p and p.get('operation') == 'PROJECT':
        return p
    if p and 'source' in p:
        return find_project_node(p.get('source'))
    return None


def extract_header(project_plan, catalog_manager):
    header = []
    is_star = project_plan['columns'] and project_plan['columns'][0]['type'] == 'star'
    if is_star:
        source_plan = project_plan.get('source', {})
        if source_plan.get('operation') == 'AGGREGATE':
            header = source_plan.get('output_schema', [])
        else:
            table_name_node = source_plan
            while 'source' in table_name_node:
                table_name_node = table_name_node['source']
            table_name = table_name_node.get('table_name')
            if table_name:
                schema = catalog_manager.get_table_schema(table_name)
                if schema:
                    header = schema.column_order
    else:
        for col in project_plan['columns']:
            if col.get('type') == 'alias':
                header.append(col['alias']['name'])
            elif col.get('type') == 'identifier':
                header.append(col['name'])
    return header


def main():
    # 构造到项目根目录的绝对路径
    project_root = os.path.abspath(os.path.dirname(__file__))
    db_file = os.path.join(project_root, "mydb.db")
    print(f"--- [Main] 正在连接数据库文件: {db_file} ---")

    file_manager = FileManager(db_file)
    buffer_manager = BufferManager(pool_size=10, file_manager=file_manager,
                                   replacement_policy_class=LRUReplacementPolicy)
    catalog_manager = CatalogManager(db_file)
    storage_engine = StorageEngine(buffer_manager, catalog_manager)
    compiler = Compiler(catalog=catalog_manager._catalog)
    executor = Executor(catalog_manager=catalog_manager, storage_engine=storage_engine, compiler=compiler)

    log_file_path = "command_log.txt"

    print("\n欢迎使用 MyDBMS!")
    print("请输入以分号 (;) 结尾的 SQL 语句。")
    print("输入 'exit' 或 'quit' 退出。")

    while True:
        try:
            sql_lines = []
            first_line = input("\nSQL> ").strip()

            if not first_line:
                continue

            cleaned_first_line = first_line.rstrip(';')
            if cleaned_first_line.lower() in ['exit', 'quit']:
                print("正在退出 MyDBMS。再见!")
                break

            if first_line.lower().startswith('source '):
                file_path_raw = first_line.split(' ', 1)[1].strip().rstrip(';')

                if (file_path_raw.startswith("'") and file_path_raw.endswith("'")) or \
                        (file_path_raw.startswith('"') and file_path_raw.endswith('"')):
                    file_path = file_path_raw[1:-1]
                else:
                    file_path = file_path_raw

                try:
                    with open(file_path, 'r', encoding='utf-8') as f:
                        script_content = f.read()

                    script_content_no_comments = re.sub(r'--.*', '', script_content)
                    sql_commands = [cmd.strip() for cmd in script_content_no_comments.split(';') if cmd.strip()]

                    print(f"在文件 '{file_path}' 中找到 {len(sql_commands)} 条有效命令。")

                    success_count = 0  # 1. 初始化成功计数器

                    # --- 3. 添加 source 命令计时开始 ---
                    source_start_time = time.time()

                    for command in sql_commands:
                        try:
                            print(f"\n> 正在执行: {command};")
                            # 向上抛出异常
                            execute_query(command + ';', compiler, executor, catalog_manager, log_file_path)

                            success_count += 1  # 2. 如果 execute_query 没报错，则计数器加一


                        except Exception as e:
                            print(f"[错误] 执行命令 '{command[:50]}...' 时失败: {e}")

                    # --- 添加 source 命令计时结束 ---
                    source_end_time = time.time()
                    source_duration = source_end_time - source_start_time

                    # 3. 打印包含成功计数的总结信息
                    total_commands = len(sql_commands)
                    print(
                        f"\n文件 '{file_path}' 中的命令执行完毕。共 {total_commands} 条命令，成功执行 {success_count} 条。")
                    print(f"--- sql文件总耗时 {source_duration:.4f} 秒 ---")


                except FileNotFoundError:
                    print(f"[错误] 脚本文件未找到: {file_path}")
                except Exception as e:
                    print(f"[错误] 读取或解析脚本时发生严重错误: {e}")
                continue

            sql_lines.append(first_line)
            current_line_input = first_line
            while not current_line_input.strip().endswith(';'):
                next_line = input("... > ").strip()
                if next_line:
                    sql_lines.append(next_line)
                    current_line_input = next_line
                if current_line_input.strip().endswith(';'):
                    break

            full_sql_query = "\n".join(sql_lines)
            execute_query(full_sql_query, compiler, executor, catalog_manager, log_file_path)

        except (RuntimeError, ValueError, SemanticError) as e:
            # 错误信息已经在 execute_query 中打印，这里我们什么都不做，
            # 循环会自然进入下一次迭代，等待用户的下一条命令。
            pass

        except (EOFError, KeyboardInterrupt):
            print("\n正在退出 MyDBMS。再见!")
            break

    print("\n正在关闭...")
    catalog_manager.save_catalog()
    buffer_manager.flush_all_pages()
    file_manager.close()
    print("关闭完成。")

if __name__ == "__main__":
    main()